import pandas as pd
import os
import csv


# import connect_mysql as cm


def get_item_info(input_file):
    """
    得到Item的信息
    input_file: Item的文件地址
    return:
        dict：  {itemID:[item_info]}
    """
    item_info = {}
    if not os.path.exists(input_file):
        return {}

    with open(input_file, "r", encoding='utf-8') as file:
        lines = csv.reader(file)
        i = 0
        for line in lines:  # 遍历每一条信息
            if i == 0:  # 跳过表头
                i += 1
                continue
            else:
                item_info[line[0]] = line[1:]
    return item_info


def get_average_score(input_file):
    """
    得到Item的平均得分
    input_file: Item的打分文件 ratings.csv
    return:
        dict  {ItemID：average_score}
    """
    score_dict = {}
    # if not os.path.exists(input_file):
    #     return {}
    # ratings_data = pd.read_csv(input_file)
    # ratings_data = cm.getdata()
    ratings_data = pd.DataFrame(list(input_file), columns=['userId', 'gameId', 'rating'])
    ratings_data[['rating']] = ratings_data[['rating']].astype(float)
    ratings_data[["gameId"]] = ratings_data[["gameId"]].astype(float)

    ratings_mean_score = ratings_data[["gameId", "rating"]].groupby("gameId").agg("mean")  # 对item分组求均值
    #movieId修改为gameId
    gameId = ratings_mean_score.index.values.astype("str")  # 将itemID 转化为str型
    mean_score = ratings_mean_score["rating"].values.round(3)  # 将均值保留三位小数

    gameId_mean_score_zip = zip(gameId, mean_score)
    for gameId, score in gameId_mean_score_zip:
        score_dict[gameId] = score
    return score_dict


def get_train_data_from_Queryset(object):
    """
        从queryset得到LFM的训练数据
        object获取query 文件
        return:
            list  [(userID, itemID, label), (userID, itemID, label)]
        """
    # if not os.path.exists(input_file):
    #     return []
    score_dict = get_average_score(object)  # item的平均得分，这里还是获取了全部数据
    pos_dict, neg_dict = {}, {}  # 正样本, 负样本
    train_data = []  # 训练集
    threshold = 3.0  # 阈值 （大于该值，为正样本；否则为负样本）

    # with open(input_file, "r", encoding='utf-8') as file:
    #     lines = csv.reader(file)
    #     i = 0
    #     for line in lines:
    #         if i == 0:  # 跳过表头
    #             i += 1
    #             continue
    #         userID, itemID, rating = line[0], line[1], float(line[2])
    #         if userID not in pos_dict:
    #             pos_dict[userID] = []
    #         if userID not in neg_dict:
    #             neg_dict[userID] = []
    #         if rating > threshold:  # rating 大于 4.0，正样本；添加到正样本中的
    #             pos_dict[userID].append((itemID, 1))
    #         else:
    #             score = score_dict.get(itemID, 0)  # 否则，获取该item 对应的平均得分；添加到负样本中
    #             neg_dict[userID].append((itemID, score))

    for line in object:
        userID, itemID, rating = line[0], line[1], float(line[2])
        if userID not in pos_dict:
            pos_dict[userID] = []
        if userID not in neg_dict:
            neg_dict[userID] = []
        if rating > threshold:  # rating 大于 4.0，正样本；添加到正样本中的
            pos_dict[userID].append((itemID, 1))
        else:
            score = score_dict.get(itemID, 0)  # 否则，获取该item 对应的平均得分；添加到负样本中
            neg_dict[userID].append((itemID, score))
    # 均衡正负样本
    for userID in pos_dict:
        data_num = min(len(pos_dict[userID]), len(neg_dict.get(userID, [])))  # 对于某用户，取其正负样本最小的数量为最终正负样本的数量
        if data_num > 0:
            train_data += [(userID, pos_data[0], pos_data[1]) for pos_data in pos_dict[userID]][
                          : data_num]  # 正样本取data_num个
        else:
            continue
        sorted_neg_list = sorted(neg_dict[userID], key=lambda x: x[1], reverse=True)[
                          : data_num]  # 根据评分对负样本排序，取前data_num个为负样本！倒序
        train_data += [(userID, neg_data[0], 0) for neg_data in sorted_neg_list]
    return train_data


def get_train_data(input_file):
    """
    得到LFM的训练数据
    input_file: user、item rating 文件
    return:
        list  [(userID, itemID, label), (userID, itemID, label)]
    """
    if not os.path.exists(input_file):
        return []
    score_dict = get_average_score(input_file)  # item的平均得分，这里还是获取了全部数据
    pos_dict, neg_dict = {}, {}  # 正样本, 负样本
    train_data = []  # 训练集
    threshold = 4.0  # 阈值 （大于该值，为正样本；否则为负样本）

    with open(input_file, "r", encoding='utf-8') as file:
        lines = csv.reader(file)
        i = 0
        for line in lines:
            if i == 0:  # 跳过表头
                i += 1
                continue
            userID, itemID, rating = line[0], line[1], float(line[2])
            if userID not in pos_dict:
                pos_dict[userID] = []
            if userID not in neg_dict:
                neg_dict[userID] = []
            if rating > threshold:  # rating 大于 4.0，正样本；添加到正样本中的
                pos_dict[userID].append((itemID, 1))
            else:
                score = score_dict.get(itemID, 0)  # 否则，获取该item 对应的平均得分；添加到负样本中
                neg_dict[userID].append((itemID, score))

    # 均衡正负样本
    for userID in pos_dict:
        data_num = min(len(pos_dict[userID]), len(neg_dict.get(userID, [])))  # 对于某用户，取其正负样本最小的数量为最终正负样本的数量
        if data_num > 0:
            train_data += [(userID, pos_data[0], pos_data[1]) for pos_data in pos_dict[userID]][
                          : data_num]  # 正样本取data_num个
        else:
            continue
        sorted_neg_list = sorted(neg_dict[userID], key=lambda x: x[1], reverse=True)[
                          : data_num]  # 根据评分对负样本排序，取前data_num个为负样本！倒序
        train_data += [(userID, neg_data[0], 0) for neg_data in sorted_neg_list]
    return train_data


if __name__ == '__main__':
    input_file = "./data/ratings.csv"  # 评分表
    train_data = get_train_data(input_file)
    print(train_data[:10])
